from abc import ABC

import torch
from torch import nn
import torch.nn.functional as F
from src.net.loss import nll_loss

from src.utils.build_model import build_encoder, build_decoder


class LASBase(nn.Module):
    def __init__(self, config):
        super(LASBase, self).__init__()
        self.config = config
        self.encoder = build_encoder(config)
        self.decoder = build_decoder(config)

    def forward(self, enc_X, enc_len, tokens, tokens_length):
        raise NotImplementedError

    def recognize(self, inputs, inputs_length):
        raise NotImplementedError

    def recognize_beam(self, inputs, inputs_length, beam=4, nbest=3):
        raise NotImplementedError


class LAS(LASBase):
    def __init__(self, config):
        super(LAS, self).__init__(config)
        self.nll_loss = nll_loss

    def forward(self, enc_X, enc_len, tokens, tokens_length):
        enc_outputs, encode_output_lengths = self.encoder(enc_X, enc_len)

        tokens_with_bos, token_with_bos_lens = F.pad(tokens, pad=[1, 0, 0, 0], value=0), tokens_length.add(1)
        tokens_with_eos, token_with_eos_lens = F.pad(tokens, pad=[0, 1, 0, 0], value=0), tokens_length.add(1)

        dec_output = self.decoder(tokens_with_bos, enc_outputs, encode_output_lengths)
        dec_output = torch.nn.functional.log_softmax(dec_output, dim=-1)
        # print("dec_output:", dec_output.shape)
        # print("tokens_with_eos:", tokens_with_eos)
        # print("token_with_eos_lens:", token_with_eos_lens)

        return self.nll_loss(dec_output, tokens_with_eos,
                             length=token_with_eos_lens)

    def recognize(self, inputs, inputs_length):
        enc_outputs, encode_output_lengths = self.encoder(inputs, inputs_length)

        output_seq = []
        for i in range(inputs.shape[0]):
            out = self.decoder.greedy_search(inputs[i],
                                             enc_outputs[i].unsqueeze(0),
                                             encode_output_lengths[i].unsqueeze(0))
            output_seq.append(out)
        # shape ( b * len)
        # print("output_seq:", output_seq)

        return output_seq

    def recognize_beam(self, inputs, inputs_length, beam=4, nbest=3):
        enc_outputs, encode_output_lengths = self.encoder(inputs, inputs_length)

        output_seq = []
        for i in range(inputs.shape[0]):
            out = self.decoder.beam_search(inputs[i],
                                           enc_outputs[i].unsqueeze(0),
                                           encode_output_lengths[i].unsqueeze(0),
                                           beam=beam, nbest=nbest)
            output_seq.append(out)
        # shape ( b * len)
        # print("output_seq:", output_seq)

        return output_seq
